package org.lim.aiagent.app;

import jakarta.annotation.Resource;
import org.lim.aiagent.advisor.ForbiddenWordsAdvisor;
import org.lim.aiagent.rag.preretrieval.QueryRewriter;
import org.lim.aiagent.tools.toolregister.ToolRegistration;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.vectorstore.VectorStore;
import lombok.extern.slf4j.Slf4j;
import org.lim.aiagent.advisor.MyLoggerAdvisor;
import org.lim.aiagent.advisor.ReReadingAdvisor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.List;

@Component
@Slf4j
public class PsychologicalApp {

    private final ChatClient chatClient;


    private final String SYSTEM_PROMPT;

    /**
     * 初始化 ChatClient
     *
     * @param dashscopeChatModel
     */
    public PsychologicalApp(ChatModel dashscopeChatModel,
                   MyLoggerAdvisor myLoggerAdvisor,
                   ReReadingAdvisor reReadingAdvisor,
                   ForbiddenWordsAdvisor forbiddenWordsAdvisor,
                   @Value("classpath:/prompts/system-message.st") org.springframework.core.io.Resource systemResource) {
        // 初始化基于文件的对话记忆
        this.SYSTEM_PROMPT = new SystemPromptTemplate(systemResource).render();
//        String fileDir = System.getProperty("user.dir") + "/tmp/chat-memory";
//        ChatMemory chatMemory = new FileBasedChatMemory(FILE_SAVE_DIR);
//         初始化基于内存的对话记忆
        MessageWindowChatMemory chatMemory = MessageWindowChatMemory.builder()
                .chatMemoryRepository(new InMemoryChatMemoryRepository())
                .maxMessages(20)
                .build();
        chatClient = ChatClient.builder(dashscopeChatModel)
                .defaultSystem(SYSTEM_PROMPT)
                .defaultAdvisors(
                        MessageChatMemoryAdvisor.builder(chatMemory).build(),
                        //日志advisor
                        myLoggerAdvisor,
                        //重读advisor
                        reReadingAdvisor,
                        //禁止词advisor
                        forbiddenWordsAdvisor
                )

                .build();
    }

    /**
     * AI 基础对话（支持多轮对话记忆）
     *
     * @param message
     * @param chatId
     * @return
     */
    public String doChat(String message, String chatId) {
        String content = null;
        try {
            ChatResponse chatResponse = chatClient
                    .prompt()
                    .user(message)
                    .advisors(spec -> spec.param(ChatMemory.CONVERSATION_ID, chatId))
                    .call()
                    .chatResponse();
            content = chatResponse.getResult().getOutput().getText();
            log.info("content: {}", content);
        } catch (IllegalArgumentException e) {
            System.err.println(e.getMessage());
            return "您的请求因包含不当内容而被拒绝。";
        }
        return content;
    }

    /**
     * AI 基础对话（支持多轮对话记忆，SSE 流式传输）
     *
     * @param message
     * @param chatId
     * @return
     */
    public Flux<String> doChatByStream(String message, String chatId) {
        return chatClient.prompt()
                .user(message)
                .advisors(spec -> spec.param(ChatMemory.CONVERSATION_ID, chatId))
                .stream()
                .content();
    }

    public record PsychologicalReport(String tittle, List<String> suggestions) {
    }

    /**
     * AI 心理心理报告功能（实战结构化输出）
     *
     * @param message
     * @param chatId
     * @return
     */
    public PsychologicalReport doChatWithReport(String message, String chatId) {
        return chatClient.prompt()
                .system(SYSTEM_PROMPT + "每次对话后都要生成心理健康报告结果")
                .user(message)
                .advisors(spec -> spec.param(ChatMemory.CONVERSATION_ID, chatId))
                .call()
                .entity(PsychologicalReport.class);
    }
    @Resource
    private VectorStore psychologicalAppVectorStore;
    @Resource
    private RetrievalAugmentationAdvisor psychologicalAppRagCloudAdvisor;
    @Resource
    private QueryRewriter queryRewriter;
    public String doChatWithRag(String message, String chatId){
        String rewrittenMessage = queryRewriter.doQueryRewrite(message);
        ChatResponse chatResponse = chatClient.prompt()
                .system(SYSTEM_PROMPT )
                //使用改写后的message
                .user(rewrittenMessage)
//                .advisors(new MyLoggerAdvisor())
                .advisors(spec -> spec.param(ChatMemory.CONVERSATION_ID, chatId))
                // 应用本地Rag知识库问答
                .advisors(new QuestionAnswerAdvisor(psychologicalAppVectorStore))
                // 应用 RAG 检索增强服务（基于云知识库服务）
                .advisors(psychologicalAppRagCloudAdvisor)
                .call()
                .chatResponse();
        assert chatResponse != null;
        String content = chatResponse.getResult().getOutput().getText();
        log.info("content: {}", content);
        return content;
    }

    // AI 调用工具能力
    @Resource
    private ToolCallback[] allTools;

    /**
     * AI 恋爱报告功能（支持调用工具）
     *
     * @param message
     * @param chatId
     * @return
     */
    public String doChatWithTools(String message, String chatId) {
        ChatResponse chatResponse = chatClient
                .prompt()
                .user(message)
                .advisors(spec -> spec.param(ChatMemory.CONVERSATION_ID, chatId))
                // 开启日志，便于观察效果
                .advisors(new MyLoggerAdvisor())
                .toolCallbacks(allTools)
                .call()
                .chatResponse();
        String content = chatResponse.getResult().getOutput().getText();
        log.info("content: {}", content);
        return content;
    }

    // AI 调用 MCP 服务

    @Resource
    private ToolCallbackProvider toolCallbackProvider;

    /**
     * AI 心理报告功能（调用 MCP 服务）
     *
     * @param message
     * @param chatId
     * @return
     */
    public String doChatWithMcp(String message, String chatId) {
        ChatResponse chatResponse = chatClient
                .prompt()
                .user(message)
                .advisors(spec -> spec.param(ChatMemory.CONVERSATION_ID, chatId))
                // 开启日志，便于观察效果
                .advisors(new MyLoggerAdvisor())
                .toolCallbacks(toolCallbackProvider)
                .call()
                .chatResponse();
        String content = chatResponse.getResult().getOutput().getText();
        log.info("content: {}", content);

        return content;
    }
}



